import python_jwt as jwt, jwcrypto.jwk as jwk, datetime

key = jwk.JWK.generate(kty='RSA', size=1024)


def authenticated_super(re):
    def _authenticated(*args, **kwargs):
        if len(args) >= 1:
            headers = dict(args[1].headers)
            token = headers.get("Authorization")
            if not token:
                return args[0].response_json(29)
            else:
                user = JWT.verify_jwt(token)
                if user:
                    if user.get("user_type") in (0, 1):
                        return re(*args, **kwargs)
                    else:
                        return args[0].response_json(111)
                else:
                    return args[0].response_json(30)
        return re(*args, **kwargs)

    return _authenticated


def authenticated(re):
    def _authenticated(*args, **kwargs):
        if len(args) >= 1:
            headers = dict(args[1].headers)
            token = headers.get("Authorization")
            if not token:
                return args[0].response_json(29)
            else:
                user = JWT.verify_jwt(token)
                if user:
                    return re(*args, **kwargs)
                else:
                    return args[0].response_json(30)
        return re(*args, **kwargs)

    return _authenticated


class JWT:

    @staticmethod
    def get_token(payload: dict) -> str:
        return jwt.generate_jwt(payload, key, 'PS256', datetime.timedelta(days=5))

    @staticmethod
    def verify_jwt(token):
        try:
            _, claims = jwt.verify_jwt(token, key, ['PS256'])
            return claims
        except Exception as e:
            print("error", e)
            return None


if __name__ == '__main__':
    token = JWT.get_token({"name": "陈毫"})
    print(token)
    # token += "dddd"
    # print(token)
    info = JWT.verify_jwt(token)
    print(info)
